import komand
from .schema import ConnectionSchema
# Custom imports below
from komand_malwareconfig.util.client import MalwareConfigClient


class Connection(komand.Connection):

    __SUCCESS_STATUS_CODE = "200"

    def __init__(self):
        super(self.__class__, self).__init__(input=ConnectionSchema())
        self.client = None

    def connect(self, params={}):
        api_key = params.get("key").get("secretKey")
        self.client = MalwareConfigClient(api_key=api_key, action=self, debug=True)

        test_connection = self.client.test_client_connection()
        response = str(test_connection.get("response"))
        self.logger.info("Connect: Response code is: {response}".format(response=response))

        # Check response code to see if we can connect to MalwareConfig
        if not response == self.__SUCCESS_STATUS_CODE:
            raise Exception("Connect: Test connection to MalwareConfig API not successful!")

        self.logger.info("Connect: Test connection to MalwareConfig API successful!")
